#include<stdio.h>
#include<stdlib.h>
#include<time.h>
#include<math.h>
#include<assert.h>

typedef struct kmeansnode{
    int id;
    float* v;
}* KMeansNode;

// calc Euclid Distance between 2 vecs
float EuclidDistance(float* v1,float* v2,int dimmension){
    float dist=0.0f;
    int i;
    for(i=0;i<dimmension;i++){
        dist+=(v1[i]-v2[i])*(v1[i]-v2[i]);
    }
    return (float)(sqrt(dist));
}

typedef struct cluster{
    int num;
    int* nodeIDs;
}* Cluster;

int sameClusters(Cluster* clusters1,Cluster* clusters2,int K){
    int k1,k2,i1,i2;
    for(k1=0;k1<K;k1++){
        Cluster c1=clusters1[k1];
        int ok=0;
        for(k2=0;k2<K;k2++){
            Cluster c2=clusters2[k2];
            if(c1->num!=c2->num) continue;
            
            int found=1;
            for(i1=0;i1<c1->num;i1++){
                int nodeID1=c1->nodeIDs[i1];
                for(i2=0;i2<c2->num;i2++){
                    if(nodeID1==c2->nodeIDs[i2]){nodeID1=-1; break;}
                    else continue;
                }
                if(nodeID1!=-1){found=0; break;}
            }
            if(found){ok=1; break;}
            else continue;
        }
        if(!ok) return 0;
        else continue;
    }
    return 1;
}

void freeClusters(Cluster* clusters,int K){
    if(clusters!=NULL){
        int k;
        for(k=0;k<K;k++){
            free(clusters[k]->nodeIDs);
            free(clusters[k]);
        }
        free(clusters);
        clusters=NULL;
    }
}

Cluster* _KMeansClustering(KMeansNode* nodes,int nodeDimension,int nodeNum,KMeansNode* teachers,int K,Cluster* prevClusters){
    int i,j,k;
    
    int nearestTeacherIDMemo[nodeNum];
    float nearestTeacherDistMemo[nodeNum];

    for(i=0;i<nodeNum;i++) nearestTeacherIDMemo[i]=-1;

    //for each node...
    for(i=0;i<nodeNum;i++){
        KMeansNode Xi=nodes[i];
        
        //for each teacher...
        for(k=0;k<K;k++){
            //calc distance
            KMeansNode teacher_k=teachers[k];
            float dist=EuclidDistance(Xi->v,teacher_k->v,nodeDimension);
            //memo the information of nearest teacher
            if(nearestTeacherIDMemo[i]<0 || nearestTeacherDistMemo[i]>dist){
                nearestTeacherIDMemo[i]=k;
                nearestTeacherDistMemo[i]=dist;
            }
        }
    }
    
    //↓ currentClusterのアローケーションと代入
    Cluster* currentClusters=(Cluster*)malloc(sizeof(Cluster)*K); assert(currentClusters!=NULL);
    for(k=0;k<K;k++){
        currentClusters[k]=(Cluster)malloc(sizeof(struct cluster)); assert(currentClusters[k]!=NULL);
        currentClusters[k]->num=0;
    }
    for(i=0;i<nodeNum;i++){
        k=nearestTeacherIDMemo[i];
        currentClusters[k]->num++;
    }
    int clustercounts[K];
    for(k=0;k<K;k++){
        currentClusters[k]->nodeIDs=(int*)malloc(sizeof(int)*currentClusters[k]->num); assert(currentClusters[k]->nodeIDs!=NULL);
        clustercounts[k]=0;
    }
    for(i=0;i<nodeNum;i++){
        k=nearestTeacherIDMemo[i];
        currentClusters[k]->nodeIDs[clustercounts[k]++]=i;
    }
    //↑ currentClusterのアローケーションと代入

    if(prevClusters!=NULL && sameClusters(prevClusters,currentClusters,K)){
        //free memory for clusters
        freeClusters(prevClusters,K);
        return currentClusters;
    }
    
    if(prevClusters!=NULL) freeClusters(prevClusters,K);
    prevClusters=currentClusters;
    currentClusters=NULL;
    
    //for each teacher...
    for(k=0;k<K;k++){
        //re-calc the center of cluster
        int n=prevClusters[k]->num;
        for(j=0;j<nodeDimension;j++) teachers[k]->v[j]=0;
        for(i=0;i<n;i++){
            for(j=0;j<nodeDimension;j++) teachers[k]->v[j]+=nodes[prevClusters[k]->nodeIDs[i]]->v[j];
        }
        for(j=0;j<nodeDimension;j++) teachers[k]->v[j]/=n;
    }
    
    return _KMeansClustering(nodes,nodeDimension,nodeNum,teachers,K,prevClusters);
}

Cluster* KMeansClustering(KMeansNode* nodes,int nodeDimension,int nodeNum,KMeansNode* teachers,int K){
    return _KMeansClustering(nodes,nodeDimension,nodeNum,teachers,K,NULL);
}

void copyNode(KMeansNode* nodes,KMeansNode* teachers,int ki,int ni,int nodeDimension){
    int i;
    teachers[ki]->id=nodes[ni]->id;
    for(i=0;i<nodeDimension;i++) teachers[ki]->v[i]=nodes[ni]->v[i];
}
float min(float a,float b){return (a<b)? a : b;}

KMeansNode* makeTeacherVectors(KMeansNode* nodes,int nodeDimension,int nodeNum,int K){
    int i,j,k,idx;
    float dist[nodeNum];
    srand((unsigned)time(NULL));

    KMeansNode* teachers=(KMeansNode*)malloc(sizeof(KMeansNode)*K); assert(teachers!=NULL);
    for(k=0;k<K;k++){
        teachers[k]=(KMeansNode)malloc(sizeof(struct kmeansnode)); assert(teachers[k]!=NULL);
        teachers[k]->v=(float*)malloc(sizeof(float)*nodeDimension); assert(teachers[k]->v!=NULL);
    }
    
    //1点目は適当に選ぶ
    k=0;
    i=rand()%nodeNum;
    copyNode(nodes,teachers,k++,i,nodeDimension);
    for(j=0;j<nodeNum;j++) dist[j]=EuclidDistance(nodes[j]->v,nodes[i]->v,nodeDimension);
    
    while(k<K){
        int bestIdx=-1;
        float bestSum;
        
        //idx=argmin_{i}(sum_{j}(min(dist[j],vec[j]-vec[i]^2))
        for(i=0;i<nodeNum;i++){
            float sum=0;
            for(j=0;j<nodeNum;j++) sum+=min(dist[j],EuclidDistance(nodes[j]->v,nodes[i]->v,nodeDimension));
            if(bestIdx<0 || sum<bestSum){bestIdx=i; bestSum=sum;}
        }
        copyNode(nodes,teachers,k++,bestIdx,nodeDimension);
        for(j=0;j<nodeNum;j++) dist[j]=min(dist[j],EuclidDistance(nodes[j]->v,nodes[bestIdx]->v,nodeDimension));
    }

    return teachers;
}

Cluster* KMeansPPClustering(KMeansNode* nodes,int nodeDimension,int nodeNum,int K){
    int k;
    KMeansNode* teachers=makeTeacherVectors(nodes,nodeDimension,nodeNum,K);
    Cluster* ret=_KMeansClustering(nodes,nodeDimension,nodeNum,teachers,K,NULL);

    for(k=0;k<K;k++){free(teachers[k]->v);free(teachers[k]);}
    free(teachers);

    return ret;
}
float bic(float nlogVar,int nodeDimension,int nodeNum,int K){
    //return -nodeDimension*nlogVar/2 + K*log(nodeNum);
    return K*log(nodeNum)+nodeNum*log(nlogVar/nodeNum);
}
Cluster* _XMeansClustering(KMeansNode* nodes,int nodeDimension,int nodeNum,int* initK){
    int i,j,k,K;
    
    //初期クラスタ。Kこに分ける。
    K=*initK;
    KMeansNode* teachers=makeTeacherVectors(nodes,nodeDimension,nodeNum,K);
    Cluster* clusters=_KMeansClustering(nodes,nodeDimension,nodeNum,teachers,K,NULL);
    float* vars=(float*)malloc(sizeof(float)*K); assert(vars!=NULL);
    
    float logVar=0;
    for(k=0;k<K;k++){
        if(clusters[k]->num<=1){
            vars[k]=0;
        }
        else{
            float* u=teachers[k]->v;
            float sqsum=0;
            for(i=0;i<clusters[k]->num;i++){
                float* v=nodes[clusters[k]->nodeIDs[i]]->v;
                for(j=0;j<nodeDimension;j++) sqsum+=(v[j]-u[j])*(v[j]-u[j]);
            }
            vars[k]=sqsum;//残差平方和
            logVar+=vars[k];
            //vars[k]=sqsum/(nodeDimension*(clusters[k]->num-1));
            //logVar+=clusters[k]->num*log(vars[k]);
        }
    }
    
    //情報量基準
    float ic=bic(logVar,nodeDimension,nodeNum,K);

    char nextOp='+',//次に行うオペレーション
    fin='n';//nextOpによって改善しないとき、終了するかとうか
    
    //while(1){
    int aaa;    
    for(aaa=0;aaa<1000;aaa++){//無限ループの対策のため、1000回までに制限
        //Kを大きくしようかなー？
        if(nextOp=='+'){
            //分散が最大のクラスタを求める
            int worstIdx=0;
            for(k=0;k<K;k++){
                if(vars[k] >= vars[worstIdx]) worstIdx=k;
            }
            //そのクラスタを2つに分ける
            KMeansNode* _nodes=(KMeansNode*)malloc(sizeof(KMeansNode)*clusters[worstIdx]->num); assert(_nodes!=NULL);
            for(i=0;i<clusters[worstIdx]->num;i++){
                _nodes[i]=(KMeansNode)malloc(sizeof(struct kmeansnode)); assert(_nodes[i]!=NULL);
                _nodes[i]->id=clusters[worstIdx]->nodeIDs[i];
                _nodes[i]->v=nodes[clusters[worstIdx]->nodeIDs[i]]->v;
            }
            KMeansNode* _teachers=makeTeacherVectors(_nodes,nodeDimension,clusters[worstIdx]->num,2);
            Cluster* _clusters=_KMeansClustering(_nodes,nodeDimension,clusters[worstIdx]->num,_teachers,2,NULL);
            for(k=0;k<2;k++){//すごく紛らわしいので、もとのインデックスに書き換え
                for(i=0;i<_clusters[k]->num;i++) _clusters[k]->nodeIDs[i]=_nodes[_clusters[k]->nodeIDs[i]]->id;
            }
            float _vars[2];
            float _logVar=0;
            for(k=0;k<2;k++){
                if(_clusters[k]->num<=1){
                    _vars[k]=0; //_logVar+=0;
                }
                else{
                    float* u=_teachers[k]->v;
                    float sqsum=0;
                    for(i=0;i<_clusters[k]->num;i++){
                        float* v=nodes[_clusters[k]->nodeIDs[i]]->v;
                        for(j=0;j<nodeDimension;j++) sqsum+=(v[j]-u[j])*(v[j]-u[j]);
                    }
                    _vars[k]=sqsum;//残差平方和
                    _logVar+=_vars[k];
                    //_vars[k]=sqsum/(_clusters[k]->num-1);
                    //_logVar+=_clusters[k]->num*log(_vars[k]);
                }
            }
            
            int _n=clusters[worstIdx]->num;
            
            //float _ic=bic(logVar-_n*log(vars[worstIdx])+_logVar,nodeDimension,nodeNum,K+1);
            float _ic=bic(logVar-vars[worstIdx]+_logVar,nodeDimension,nodeNum,K+1);
            if(_ic<ic){

//printf("[%c] %f %f\n",nextOp,ic,_ic);
//printf("[ "); for(i=0;i<_n;i++) printf("%d ",clusters[worstIdx]->nodeIDs[i]); printf("]\n");
//for(k=0;k<2;k++){
//    printf("=>[ "); for(i=0;i<_clusters[k]->num;i++) printf("%d ",_clusters[k]->nodeIDs[i]); printf("]\n");
//}
//printf("\n");
                //分ける
                clusters=(Cluster*)realloc(clusters,sizeof(Cluster)*(K+1)); assert(clusters!=NULL);
                teachers=(KMeansNode*)realloc(teachers,sizeof(KMeansNode)*(K+1)); assert(teachers!=NULL);
                vars=(float*)realloc(vars,sizeof(float)*(K+1)); assert(vars!=NULL);
                
                clusters[K]=(Cluster)malloc(sizeof(struct cluster)); assert(clusters[K]!=NULL);
                clusters[K]->num=_clusters[0]->num;
                clusters[K]->nodeIDs=_clusters[0]->nodeIDs;
                free(_clusters[0]);
    

                teachers[K]=(KMeansNode)malloc(sizeof(struct kmeansnode)); assert(teachers[K]!=NULL);
                teachers[K]->v=_teachers[0]->v;
                free(_teachers[0]);
                
                clusters[worstIdx]->num=_clusters[1]->num;
                free(clusters[worstIdx]->nodeIDs);
                clusters[worstIdx]->nodeIDs=_clusters[1]->nodeIDs;
                free(_clusters[1]);
                
                free(teachers[worstIdx]->v);
                teachers[worstIdx]->v=_teachers[1]->v;
                free(_teachers[1]);
                
                free(_teachers);
                free(_clusters);
                for(i=0;i<_n;i++) free(_nodes[i]);
                free(_nodes);
                
                //logVar+=_logVar-_n*log(vars[worstIdx]);
                logVar+=_logVar-vars[worstIdx];
                vars[K]=_vars[0]; vars[worstIdx]=_vars[1];
                ic=_ic;
                K++;
                fin='n';
            }
            else{
                //分けない
                for(k=0;k<2;k++){free(_teachers[k]->v);free(_teachers[k]);}
                free(_teachers);
                for(i=0;i<_n;i++) free(_nodes[i]);
                free(_nodes);
                freeClusters(_clusters,2);
                if(fin=='n'){ nextOp='-'; fin='y'; }
                else break;
            }
        }
        //Kを小さくしようかなー？
        else if(nextOp='-'){
            if(K<=1){
                if(fin=='n'){ nextOp='+'; fin='y'; continue;}
                else break;
            }
            
            //クラス間距離が最小の2クラスタを求める
            int bestIdx1=-1,bestIdx2;
            float bestDist;
            for(k=0;k<K-1;k++){
                for(j=k+1;j<K;j++){
                    float dist=0;
                    for(i=0;i<nodeDimension;i++) dist+=(teachers[k]->v[i]-teachers[j]->v[i])*(teachers[k]->v[i]-teachers[j]->v[i]);
                    if(bestIdx1<0 || dist<bestDist){
                        bestIdx1=k; bestIdx2=j;
                        bestDist=dist;
                    }
                }
            }
            //そのクラスタをくっつける
            float* _teacherV=(float*)malloc(sizeof(float)*nodeDimension); assert(_teacherV!=NULL);
            int n1=clusters[bestIdx1]->num,n2=clusters[bestIdx2]->num;
            int _num=n1+n2;
            for(i=0;i<nodeDimension;i++){
                _teacherV[i]=(teachers[bestIdx1]->v[i]*n1
                            +teachers[bestIdx2]->v[i]*n2)/_num;
            }
            float _logVar;
            float _var=0;
            if(n1<=1){
                _logVar=0;
            }
            else{
                for(i=0;i<n1;i++){
                    float* v=nodes[clusters[bestIdx1]->nodeIDs[i]]->v;
                    for(j=0;j<nodeDimension;j++) _var+=(v[j]-_teacherV[j])*(v[j]-_teacherV[j]);
                }
                for(i=0;i<n2;i++){
                    float* v=nodes[clusters[bestIdx2]->nodeIDs[i]]->v;
                    for(j=0;j<nodeDimension;j++) _var+=(v[j]-_teacherV[j])*(v[j]-_teacherV[j]);
                }
                _logVar=_var;
                //_var/=(nodeDimension*(_num-1));
                //_logVar=_num*log(_var);

            }

            //float _ic=bic(logVar -n1*log(vars[bestIdx1]) -n2*log(vars[bestIdx2]) +_logVar,nodeDimension,nodeNum,K-1);
            float _ic=bic(logVar -vars[bestIdx1] -vars[bestIdx2] +_logVar,nodeDimension,nodeNum,K-1);
            if(_ic<ic){
            
//printf("[%c] %f %f\n",nextOp,ic,_ic);
//printf("[ "); for(i=0;i<n1;i++) printf("%d ",clusters[bestIdx1]->nodeIDs[i]); printf("]\n");
//printf("+[ "); for(i=0;i<n2;i++) printf("%d ",clusters[bestIdx2]->nodeIDs[i]); printf("]\n");
//printf("\n");
                //くっつける
                clusters[bestIdx1]->nodeIDs=(int*)realloc(clusters[bestIdx1]->nodeIDs,sizeof(int)*(_num)); assert(clusters[bestIdx1]->nodeIDs!=NULL);
                for(i=0;i<n2;i++) clusters[bestIdx1]->nodeIDs[n1+i]=clusters[bestIdx2]->nodeIDs[i];
                clusters[bestIdx1]->num=_num;
                free(clusters[bestIdx2]->nodeIDs);
                clusters[bestIdx2]->num=clusters[K-1]->num;
                clusters[bestIdx2]->nodeIDs=clusters[K-1]->nodeIDs;
                free(clusters[K-1]);
                clusters=(Cluster*)realloc(clusters,sizeof(Cluster)*(K-1)); assert(clusters!=NULL);
                
                free(teachers[bestIdx1]->v); free(teachers[bestIdx2]->v);
                teachers[bestIdx1]->v=_teacherV;
                teachers[bestIdx2]->v=teachers[K-1]->v;
                free(teachers[K-1]);
                teachers=(KMeansNode*)realloc(teachers,sizeof(KMeansNode)*(K-1)); assert(teachers!=NULL);
                
                //logVar+=_logVar-n1*log(vars[bestIdx1])-n2*log(vars[bestIdx2]);
                logVar+=_logVar-vars[bestIdx1]-vars[bestIdx2];
                vars[bestIdx1]=_var;
                vars[bestIdx2]=vars[K-1];
                vars=(float*)realloc(vars,sizeof(float)*(K-1)); assert(vars!=NULL);
                
                ic=_ic;
                K--;
                fin='n';
            }
            else{
                //くっつけない
                free(_teacherV);
                
                if(fin=='n'){ nextOp='+'; fin='y'; }
                else break;
            }
        }
    }

    //引数の書き換え。
    *initK=K;

    for(k=0;k<K;k++){free(teachers[k]->v);free(teachers[k]);}
    free(teachers);
    free(vars);
    return clusters;
}
Cluster* XMeansClustering(KMeansNode* nodes,int nodeDimension,int nodeNum,int* initK){
    int aaa;
    int k;
    Cluster* clusters=NULL;
    for(aaa=0;aaa<5;aaa++){
        if(clusters!=NULL){ freeClusters(clusters,k); clusters=NULL;}
        k=*initK+aaa;

        clusters=_XMeansClustering(nodes,nodeDimension,nodeNum,&k);
        if(k>1) break;
        //クラスタ数が1の時はやり直し
    }
    *initK=k;
    return clusters;
}
